#用来可视化关键点坐标

import os
from pathlib import Path
import cv2

def visualize_keypoints_txt(image_dir, label_dir, output_dir, img_exts=('jpg','png','jpeg')):
    """
    从图片目录和对应的txt标签文件中读取关键点坐标，并在图像上绘制点及编号。
    假设每个txt文件与图片同名，不同后缀，
    文本格式为（YOLO格式拓展）：
      class_id x_center y_center w h [x1 y1 v1 x2 y2 v2 ...]
    其中前5个字段为bbox信息，后续字段为关键点序列，每3个一组：x, y, v(v可见性)，
    坐标支持归一化或像素值。
    同时在图像上绘制每个bbox的中心点，不同class使用不同颜色，
    并且为每个对象分别编号关键点。
    """
    image_dir = Path(image_dir)
    label_dir = Path(label_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 定义不同class_id的颜色（BGR格式）
    colors = [
        (0, 255, 0),    # class 0: 绿色
        (255, 0, 0),    # class 1: 蓝色
        (0, 0, 255),    # class 2: 红色
        (0, 255, 255),  # class 3: 黄色
        (255, 0, 255),  # class 4: 洋红
        (255, 255, 0),  # class 5: 青色
    ]

    for img_path in image_dir.iterdir():
        if not img_path.suffix.lower().lstrip('.') in img_exts:
            continue
        base = img_path.stem
        txt_path = label_dir / f"{base}.txt"
        if not txt_path.exists():
            print(f"Warning: label file not found for {img_path.name}")
            continue

        img = cv2.imread(str(img_path))
        if img is None:
            print(f"Warning: failed to load image {img_path.name}")
            continue
        h, w = img.shape[:2]

        with open(txt_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        for obj_id, line in enumerate(lines):
            parts = line.strip().split()
            if len(parts) < 5:
                continue
            # 解析bbox中心和class_id
            try:
                cls = int(parts[0])
                x_c = float(parts[1])
                y_c = float(parts[2])
                # 处理归一化坐标
                cx = int(x_c * w) if 0 <= x_c <= 1 else int(x_c)
                cy = int(y_c * h) if 0 <= y_c <= 1 else int(y_c)
            except ValueError:
                continue

            # 绘制中心点
            color = colors[cls % len(colors)]
            cv2.circle(img, (cx, cy), 5, color, -1)
            cv2.putText(img, f"C{cls}", (cx+8, cy+8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

            # 解析关键点
            kp_data = parts[5:]
            kp_points = []
            for i in range(0, len(kp_data), 3):
                try:
                    x = float(kp_data[i])
                    y = float(kp_data[i+1])
                    v = float(kp_data[i+2])
                except (IndexError, ValueError):
                    break
                if v <= 0:
                    continue
                px = int(x * w) if 0 <= x <= 1 else int(x)
                py = int(y * h) if 0 <= y <= 1 else int(y)
                kp_points.append((px, py))

            # 为当前对象绘制关键点
            for idx, (px, py) in enumerate(kp_points):
                cv2.circle(img, (px, py), 5, (0, 0, 255), -1)
                cv2.putText(img, str(idx), (px+8, py-8), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

        out_path = output_dir / img_path.name
        cv2.imwrite(str(out_path), img)
        print(f"Processed: {img_path.name}, {len(lines)} objects")

if __name__ == '__main__':
    image_dir = r"C:\Users\lenovo\Desktop\tup\images"
    label_dir = r"C:\Users\lenovo\Desktop\tup\labels"
    output_dir = r"D:\tup_yolox\visualization_results"
    visualize_keypoints_txt(image_dir, label_dir, output_dir)
